import pymysql
import json
import sys
import logging
from datetime import datetime

# Configure logging format
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)


def log_info(message):
    logging.info(message)


def log_error(message):
    logging.error(message)


# Process JSON config file
def load_config(config_file):
    try:
        with open(config_file, 'r') as f:
            config = json.load(f)
        return config
    except Exception as e:
        log_error(f"Failed to load config file: {e}")
        sys.exit(1)


# Get database connection
def get_db_connection(host, port, user, password, database):
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor,
            connect_timeout=30
        )
        return conn
    except Exception as e:
        log_error(f"Database connection failed: {e}")
        sys.exit(1)


# Get table creation order to handle foreign key dependencies
def get_creation_order(tables):
    # Simple implementation: create tables without FKs first, then those with FKs
    tables_with_fk = []
    tables_without_fk = []

    for table, info in tables.items():
        if info['type'] == 'VIEW':
            continue  # Views will be handled separately
        if "FOREIGN KEY" in info['create_sql'].upper():
            tables_with_fk.append(table)
        else:
            tables_without_fk.append(table)

    return tables_without_fk + tables_without_fk


# Get all table and view structures from source database
def get_source_objects(source_conn):
    objects = {}
    try:
        with source_conn.cursor() as cursor:
            # Get all table names
            cursor.execute("SHOW FULL TABLES WHERE Table_type = 'BASE TABLE'")
            table_names = [row[f"Tables_in_{source_conn.db.decode()}"] for row in cursor.fetchall()]

            # Get all view names
            cursor.execute("SHOW FULL TABLES WHERE Table_type = 'VIEW'")
            view_names = [row[f"Tables_in_{source_conn.db.decode()}"] for row in cursor.fetchall()]

            # Get structure for each table
            for table in table_names:
                try:
                    # Get table creation SQL
                    cursor.execute(f"SHOW CREATE TABLE `{table}`")
                    create_info = cursor.fetchone()

                    if not create_info:
                        log_error(f"Failed to get creation info for table {table}")
                        continue

                    create_table_sql = create_info['Create Table']

                    # Get table column info
                    cursor.execute(f"DESCRIBE `{table}`")
                    columns = cursor.fetchall()

                    objects[table] = {
                        'type': 'TABLE',
                        'create_sql': create_table_sql,
                        'columns': columns
                    }

                    log_info(f"Successfully got structure for table {table}")
                except Exception as e:
                    log_error(f"Error getting structure for table {table}: {e}")
                    continue

            # Get structure for each view
            for view in view_names:
                try:
                    # Get view creation SQL
                    cursor.execute(f"SHOW CREATE VIEW `{view}`")
                    create_info = cursor.fetchone()

                    if not create_info:
                        log_error(f"Failed to get creation info for view {view}")
                        continue

                    create_view_sql = create_info['Create View']

                    # Get view column info
                    cursor.execute(f"DESCRIBE `{view}`")
                    columns = cursor.fetchall()

                    objects[view] = {
                        'type': 'VIEW',
                        'create_sql': create_view_sql,
                        'columns': columns
                    }

                    log_info(f"Successfully got structure for view {view}")
                except Exception as e:
                    log_error(f"Error getting structure for view {view}: {e}")
                    continue

        return objects
    except Exception as e:
        log_error(f"Failed to get source table/view structures: {e}")
        sys.exit(1)


# Sync objects to target database
def sync_objects_to_target(source_objects, target_conn):
    try:
        with target_conn.cursor() as cursor:
            # Check what objects already exist in target
            cursor.execute("SHOW FULL TABLES")
            existing_objects = {
                row[f"Tables_in_{target_conn.db.decode()}"]: row['Table_type']
                for row in cursor.fetchall()
            }

            # First sync tables (handling dependencies)
            tables_to_sync = {
                name: info for name, info in source_objects.items()
                if info['type'] == 'TABLE'
            }

            if tables_to_sync:
                log_info("Starting table sync...")
                creation_order = get_creation_order(tables_to_sync)

                for table in creation_order:
                    if table not in tables_to_sync:
                        continue

                    table_info = tables_to_sync[table]

                    if table in existing_objects:
                        log_info(f"Table {table} already exists, skipping creation")
                        continue

                    try:
                        # Temporarily disable FK checks
                        cursor.execute("SET FOREIGN_KEY_CHECKS=0")

                        # Execute table creation
                        log_info(f"Creating table: {table}")
                        cursor.execute(table_info['create_sql'])
                        log_info(f"Table {table} created successfully")

                        # Re-enable FK checks
                        cursor.execute("SET FOREIGN_KEY_CHECKS=1")

                    except Exception as e:
                        log_error(f"Failed to create table {table}: {e}")
                        # Skip this table and continue
                        target_conn.rollback()
                        cursor.execute("SET FOREIGN_KEY_CHECKS=1")
                        continue

            # Then sync views
            views_to_sync = {
                name: info for name, info in source_objects.items()
                if info['type'] == 'VIEW'
            }

            if views_to_sync:
                log_info("Starting view sync...")
                for view, view_info in views_to_sync.items():
                    if view in existing_objects:
                        log_info(f"View {view} already exists, dropping and recreating")
                        try:
                            cursor.execute(f"DROP VIEW IF EXISTS `{view}`")
                        except Exception as e:
                            log_error(f"Failed to drop view {view}: {e}")
                            continue

                    try:
                        # Create or replace view
                        log_info(f"Creating view: {view}")
                        cursor.execute(view_info['create_sql'])
                        log_info(f"View {view} created successfully")
                    except Exception as e:
                        log_error(f"Failed to create view {view}: {e}")
                        target_conn.rollback()
                        continue

        target_conn.commit()
        log_info("All objects synced successfully")
    except Exception as e:
        target_conn.rollback()
        log_error(f"Failed to sync objects: {e}")
        sys.exit(1)


# Main function
def main(config_file):
    start_time = datetime.now()
    log_info(f"Starting schema sync, begin time: {start_time}")

    # Load config
    config = load_config(config_file)
    source_config = config['source_db']
    target_config = config['target_db']

    # Get database connections
    log_info("Connecting to source database...")
    source_conn = get_db_connection(
        source_config['host'],
        source_config['port'],
        source_config['user'],
        source_config['password'],
        source_config['database']
    )

    log_info("Connecting to target database...")
    target_conn = get_db_connection(
        target_config['host'],
        target_config['port'],
        target_config['user'],
        target_config['password'],
        target_config['database']
    )

    # Get source table and view structures
    log_info("Getting source database objects...")
    source_objects = get_source_objects(source_conn)
    log_info(
        f"Got structures for {len([o for o in source_objects.values() if o['type'] == 'TABLE'])} tables and {len([o for o in source_objects.values() if o['type'] == 'VIEW'])} views")

    # Sync to target database
    log_info("Starting schema sync to target database...")
    sync_objects_to_target(source_objects, target_conn)

    # Close connections
    source_conn.close()
    target_conn.close()

    end_time = datetime.now()
    duration = end_time - start_time
    log_info(f"Sync completed, end time: {end_time}")
    log_info(f"Total duration: {duration}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        log_error("Usage: python sync_mysql_schema.py config.json")
        sys.exit(1)

    main(sys.argv[1])
